#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
碳排放预测与优化模块

功能：
1. 碳排放数据预处理与分析
2. 机器学习模型预测碳排放量
3. 提供碳排放优化建议
4. 可视化分析结果
"""

import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error
from scipy.optimize import linprog
import matplotlib.pyplot as plt
import seaborn as sns

class CarbonEmissionSystem:
    """碳排放预测与优化核心类"""
    
    def __init__(self):
        self.model = None
        self.features = ["energy_consumption", "industrial_output", "renewable_ratio"]
        self.target = "carbon_emission"
        
    def prepare_data(self, raw_data):
        """
        数据预处理
        
        参数:
            raw_data: 原始数据 (dict或DataFrame)
            
        返回:
            处理后的数据 (DataFrame)
        """
        if isinstance(raw_data, dict):
            df = pd.DataFrame(raw_data)
        else:
            df = raw_data.copy()
            
        # 确保列名标准化
        df.columns = df.columns.str.lower()
        
        # 检查必要字段
        for col in self.features + [self.target]:
            if col not in df.columns:
                raise ValueError(f"缺少必要列: {col}")
                
        return df[self.features + [self.target]]
    
    def train_model(self, data):
        """
        训练碳排放预测模型
        
        参数:
            data: 训练数据 (DataFrame)
            
        返回:
            模型评估结果 (dict)
        """
        df = self.prepare_data(data)
        
        # 划分训练集和测试集
        X = df[self.features]
        y = df[self.target]
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=0.2, random_state=42
        )
        
        # 训练随机森林模型
        self.model = RandomForestRegressor(n_estimators=100, random_state=42)
        self.model.fit(X_train, y_train)
        
        # 评估模型
        y_pred = self.model.predict(X_test)
        mse = mean_squared_error(y_test, y_pred)
        
        return {
            "model": "RandomForestRegressor",
            "mse": mse,
            "feature_importance": dict(zip(self.features, self.model.feature_importances_))
        }
    
    def predict_emission(self, input_data):
        """
        预测碳排放量
        
        参数:
            input_data: 输入数据 (dict或DataFrame)
            
        返回:
            预测结果 (float)
        """
        if self.model is None:
            raise ValueError("模型未训练，请先调用train_model方法")
            
        df = pd.DataFrame([input_data]) if isinstance(input_data, dict) else input_data
        return self.model.predict(df[self.features])[0]
    
    def optimize_emission(self, current_data):
        """
        碳排放优化建议
        
        参数:
            current_data: 当前状态数据 (dict)
            
        返回:
            优化结果 (dict)
        """
        # 固定能源消耗和工业产值，优化新能源占比
        c = [-0.5]  # 目标函数系数：新能源占比每提高1%，碳排放减少0.5吨
        
        # 约束条件：新能源占比 <= 30
        A_ub = [[1]]
        b_ub = [30]
        
        # 边界条件
        bounds = [(0, 30)]
        
        # 求解线性规划
        result = linprog(c, A_ub=A_ub, b_ub=b_ub, bounds=bounds, method='highs')
        optimal_ratio = result.x[0]
        
        return {
            "current_emission": self.predict_emission(current_data),
            "optimal_renewable_ratio": optimal_ratio,
            "optimized_emission": self.predict_emission({
                **current_data,
                "renewable_ratio": optimal_ratio
            }),
            "reduction_potential": None  # 将在下面计算
        }
    
    def visualize(self, data, output_path=None):
        """
        数据可视化
        
        参数:
            data: 要可视化的数据 (DataFrame)
            output_path: 图片保存路径 (可选)
            
        返回:
            matplotlib Figure对象
        """
        df = self.prepare_data(data)
        
        # 创建可视化
        fig, axes = plt.subplots(2, 2, figsize=(12, 10))
        
        # 能源消耗 vs 碳排放
        sns.scatterplot(
            x="energy_consumption", 
            y="carbon_emission",
            data=df,
            ax=axes[0, 0]
        )
        axes[0, 0].set_title("能源消耗 vs 碳排放")
        
        # 工业产值 vs 碳排放
        sns.scatterplot(
            x="industrial_output", 
            y="carbon_emission",
            data=df,
            ax=axes[0, 1]
        )
        axes[0, 1].set_title("工业产值 vs 碳排放")
        
        # 新能源占比 vs 碳排放
        sns.scatterplot(
            x="renewable_ratio", 
            y="carbon_emission",
            data=df,
            ax=axes[1, 0]
        )
        axes[1, 0].set_title("新能源占比 vs 碳排放")
        
        # 特征重要性
        if self.model:
            importance = pd.DataFrame({
                "feature": self.features,
                "importance": self.model.feature_importances_
            }).sort_values("importance", ascending=False)
            
            sns.barplot(
                x="importance",
                y="feature",
                data=importance,
                ax=axes[1, 1]
            )
            axes[1, 1].set_title("特征重要性")
        
        plt.tight_layout()
        
        if output_path:
            plt.savefig(output_path)
            
        return fig

# 示例用法
if __name__ == "__main__":
    # 模拟数据
    data = {
        "energy_consumption": np.random.uniform(1000, 5000, 100),
        "industrial_output": np.random.uniform(50, 200, 100),
        "renewable_ratio": np.random.uniform(0, 30, 100),
        "carbon_emission": np.random.uniform(5000, 20000, 100)
    }
    
    # 初始化系统
    system = CarbonEmissionSystem()
    
    # 训练模型
    results = system.train_model(data)
    print(f"模型训练完成，MSE: {results['mse']:.2f}")
    print("特征重要性:", results["feature_importance"])
    
    # 预测示例
    sample_input = {
        "energy_consumption": 2500,
        "industrial_output": 120,
        "renewable_ratio": 15
    }
    prediction = system.predict_emission(sample_input)
    print(f"预测碳排放量: {prediction:.2f} 吨CO2")
    
    # 优化示例
    optimization = system.optimize_emission(sample_input)
    print(f"优化建议: 新能源占比提高到 {optimization['optimal_renewable_ratio']:.2f}%")
    print(f"预计碳排放减少: {optimization['current_emission'] - optimization['optimized_emission']:.2f} 吨CO2")
    
    # 可视化
    system.visualize(data)
    plt.show()